# Alexander F. Gazmararian
# afg2@princeton.edu

# load packages
library(tidyverse)
library(here)
library(modelsummary)
library(sandwich)
library(lmtest)

source(here("code","fun", "table_spec.R"))

# Load data
fair <- readRDS(here("data", "fair_conjoint.rds"))
fair <- subset(fair, !is.na(selected))

fair <-
  fair %>%
  rename(
    Free.Retraining.Program = `Free Retraining Program`,
    Retrained.Worker.Salary = `Retrained Worker Salary`,
    Worker.Retraining.Time = `Worker Retraining Time`,
    Income.Support.During.Retraining = `Income Support During Retraining`,
    Relocation.Support = `Relocation Support`,
    Community.Investment = `Community Investment`,
    Benefit.Support = `Benefit Support`
  )

# Specify empirical model
f <- selected ~ Community.Investment + Free.Retraining.Program + Income.Support.During.Retraining + 
  Worker.Retraining.Time + Benefit.Support + Retrained.Worker.Salary + Relocation.Support

# Function ----
est_het_conjoint <- function(intervar, intervarlab, modelvar) {
  #extract and combine coefficients
  cj <- modelplot(m, vcov = ~ ResponseId, draw = FALSE)
  #combine and estimate difference
  f.inter <- as.formula(paste0("selected ~ (Community.Investment + Free.Retraining.Program + Income.Support.During.Retraining + 
    Worker.Retraining.Time + Benefit.Support + Retrained.Worker.Salary + 
    Relocation.Support)", "*", intervar))
  
  m.int <- lm(f.inter, fair, weights = weights)
  df.int <- modelplot(m.int, vcov = ~ ResponseId, draw = FALSE)
  #add together and get p-value for difference
  df.int.merge <- subset(df.int, select = c(p.value,term), grepl(intervarlab, term))
  names(df.int.merge)[1] <- "p.df"
  df.int.merge$term <- gsub(intervarlab, "", df.int.merge$term)
  plot.df <-
    cj %>%
    left_join(., df.int.merge, by = "term") %>%
    mutate(p.df = ifelse(model== modelvar, NA, p.df),
           p.df = case_when(
             p.df >= .05 & p.df <.1 ~ "*",
             p.df < 0.05 & p.df >= .01 ~ "**",
             p.df < 0.01 ~ "***",
             T ~ NA
           )) %>%
    filter(term!="(Intercept)" & !grepl("task|ResponseId",term)) %>%
    mutate(lb90 = estimate + std.error * qnorm(.05),
           ub90 = estimate + std.error * qnorm(.95))
  plot.df <-
    plot.df %>%
    mutate(
      attribute = case_when(
        grepl("Benefit.Support", term) ~ "Benefit Support",
        grepl("Community.Investment", term) ~ "Community Investment",
        grepl("Free.Retraining", term) ~ "Free Retraining Program",
        grepl("Relocation", term) ~ "Relocation Support",
        grepl("Retrained.Worker", term) ~ "Retrained Worker Salary",
        grepl("Worker.Retraining.Time", term) ~ "Worker Retraining Time",
        grepl("Income.Support", term) ~ "Income Support During Retraining"
      )
    )
  #relevel
  baseline <- cbind(estimate = 0, std.error = 0)
  reflevels <- c(
    levels(fair$Benefit.Support)[[1]],
    levels(fair$Community.Investment)[[1]],
    levels(fair$Free.Retraining)[[1]],
    levels(fair$Relocation)[[1]],
    levels(fair$Worker.Retraining.Time)[[1]],
    levels(fair$Income.Support)[[1]]
  )
  attributelevels <- c("Benefit Support", "Community Investment", "Free Retraining Program", "Relocation Support", "Retrained Worker Salary", "Income Support During Retraining")
  #baseline dataframe
  base.df <- data.frame(term = reflevels, baseline, attribute = attributelevels)
  base.df <- mutate(base.df, across(estimate:std.error, ~ as.numeric(.x)))
  # Merge
  con <- bind_rows(plot.df, base.df)
  # Relevel factors
  con$term[con$term == "None" & con$Attribute == "Benefit Support"] <- "No Benefit Support"
  con$term[con$term == "None" & con$Attribute == "Income Support During Retraining"] <- "No Income Support"
  con$term[con$term == "None" & con$Attribute == "Relocation Support"] <- "No Relocation Support"
  con$term[con$term == "$200 per week for all community members"] <- "$200 per week for all"
  #Rename term
  con <-
    con %>%
    mutate(
      term = case_when(
        term == "Community.InvestmentBroadband" ~ "Broadband",
        term == "Community.InvestmentHousing for new residents" ~ "Housing for new residents",
        term == "Free.Retraining.ProgramFor jobs in clean energy" ~ "For jobs in clean energy",
        term == "Free.Retraining.ProgramFor jobs in health care" ~ "For jobs in health care",
        term == "Income.Support.During.Retraining$400 per week for workers" ~ "$400 per week for workers",
        term == "Income.Support.During.RetrainingNone" ~ "None",
        term == "Worker.Retraining.Time2 years" ~ "2 years",
        term == "Worker.Retraining.Time3-6 months" ~ "3-6 months",
        term == "Benefit.SupportFunding for worker health care" ~ "Funding for worker health care",
        term == "Benefit.SupportFunding for worker pensions" ~ "Funding for worker pensions",
        term == "Relocation.SupportVoucher to cover moving expenses" ~ "Voucher to cover moving expenses",
        term == "Retrained.Worker.Salary$100,000 per year" ~ "$100,000 per year",
        term == "Retrained.Worker.Salary$50,000 per year" ~ "$50,000 per year",
        T ~ NA_character_
      )
    )
  #Create data frames
  benefit <- data.frame(attribute = "Benefit Support", term = "Benefit Support")
  community <- data.frame(attribute = "Community Investment", term = "Community Investment")
  train <- data.frame(attribute = "Free Retraining Program", term = "Free Retraining Program")
  income <- data.frame(attribute = "Income Support During Retraining", term = "Income Support During Retraining")
  locate <- data.frame(attribute = "Relocation Support", term = "Relocation Support")
  salary <- data.frame(attribute = "Retrained Worker Salary", term = "Retrained Worker Salary")
  time <- data.frame(attribute = "Worker Retraining Time", term = "Worker Retraining Time")
  con <- bind_rows(con, benefit, community, train, income, locate, salary, time)
  # Arrange by level
  con <-
    con %>%
    arrange(attribute, !is.na(estimate))
  # Fix names
  con$term[[1]] <- "**Benefit Support**<br>Baseline: None"
  con$term[[7]] <- "**Community Investment**<br>Baseline: Schools"
  con$term[[13]] <- "**Free Retraining Program**<br>Baseline: For jobs in manufacturing"
  con$term[[19]] <- "**Income Support**<br>Baseline: $200 per week for all"
  con$term[[25]] <- "**Relocation Support**<br>Baseline: None"
  con$term[[29]] <- "**Retrained Worker Salary**<br>Baseline: $75,000 per year"
  con$term[[35]] <- "**Worker Retraining Time**<br>Baseline: 1 year"
  
  con <- subset(con, !is.na(term))
  
  # Create levels
  con$term <- factor(con$term, ordered = TRUE, levels = rev(c(
    "**Relocation Support**<br>Baseline: None", "Voucher to cover moving expenses",
    "**Community Investment**<br>Baseline: Schools", "Broadband", "Housing for new residents",
    "**Benefit Support**<br>Baseline: None", "Funding for worker health care", "Funding for worker pensions",
    "**Income Support**<br>Baseline: $200 per week for all", "$400 per week for workers",  "None",
    "**Free Retraining Program**<br>Baseline: For jobs in manufacturing", "For jobs in clean energy", "For jobs in health care",
    "**Retrained Worker Salary**<br>Baseline: $75,000 per year", "$100,000 per year",  "$50,000 per year",
    "**Worker Retraining Time**<br>Baseline: 1 year", "2 years", "3-6 months"
  )))
  
  back_line <- 0.25
  pos_dodge <- .7
  plot.cj <-
    con %>%
    filter(!is.na(term)) %>%
    ggplot(aes(x=term,y=estimate,color=model,shape=model)) +
    geom_vline(xintercept = "**Community Investment**<br>Baseline: Schools", color = "gray", size = back_line) +
    geom_vline(xintercept = "**Benefit Support**<br>Baseline: None", color = "gray", size = back_line) +
    geom_vline(xintercept = "**Free Retraining Program**<br>Baseline: For jobs in manufacturing", color = "gray", size = back_line) +
    geom_vline(xintercept = "**Income Support**<br>Baseline: $200 per week for all", color = "gray", size = back_line) +
    geom_vline(xintercept = "**Retrained Worker Salary**<br>Baseline: $75,000 per year", color = "gray", size = back_line) +
    geom_vline(xintercept = "**Worker Retraining Time**<br>Baseline: 1 year", color = "gray", size = back_line) +
    geom_hline(yintercept=0, color = "grey") +
    geom_errorbar(aes(ymin=conf.low,ymax=conf.high),width=0,position=position_dodge(pos_dodge)) +
    geom_errorbar(aes(ymin=lb90,ymax=ub90),width=0,size=1.25,position=position_dodge(pos_dodge)) +
    geom_point(size = 3,fill="white",position=position_dodge(pos_dodge)) +
    geom_text(aes(x=term,label=p.df), size = 6, show.legend = FALSE) +
    scale_color_grey(na.translate = FALSE, start = 0, end = .5) +
    scale_shape_manual(values = c(21,24), na.translate = FALSE) +
    theme_bw(base_size = 14) +
    coord_flip() +
    scale_y_continuous(labels = scales::percent) +
    labs(color="",shape="",y="Change in Probability of Support", x="") +
    theme(
      panel.grid = element_blank(),
      axis.text.y = ggtext::element_markdown(),
      legend.position = "bottom",
      legend.background = element_blank(),
      legend.box.background = element_rect(color = "black")
    )
  plot.cj
  return(plot.cj)
}

# Heterogeneous effects by sex ----

#impute sex
fair$sex <- with(fair, ifelse(is.na(sex), 1, sex))

# estimate models
m <- list()
m[["Female"]] <- lm(f, fair, subset = (sex == 1), weights = weights)
m[["Male"]] <- lm(f, fair, subset = (sex == 0), weights = weights)

het.sex <- est_het_conjoint(
  intervar = "sex",
  intervarlab = ":sex",
  modelvar = "Male"
)
het.sex

ggsave(
  het.sex,
  filename = here("output", "figures", "si_fig_5.1.png"),
  scale = 1.25,
  width = 6.5,
  height = 6.5,
  dpi = 300
)

# Heterogeneous effects by attentiveness ----


# estimate models
m <- list()
m[["Higher Attentiveness"]] <- lm(f, fair, subset = (pass_attn == 1), weights = weights)
m[["Lower Attentiveness"]] <- lm(f, fair, subset = (pass_attn == 0), weights = weights)

het.attn <- est_het_conjoint(
  intervar = "pass_attn",
  intervarlab = ":pass_attn",
  modelvar = "Lower Attentiveness"
)
het.attn

ggsave(
  het.attn,
  filename = here("output", "figures", "si_fig_5.2.png"),
  scale = 1.25,
  width = 6.5,
  height = 6.5,
  dpi = 300
)

# Heterogeneous effects by economic conservatism ----

# estimate models
m <- list()
m[["Economically Liberal"]] <- lm(f, fair, subset = (welfare_sup == "Support"), weights = weights)
m[["Economically Conservative"]] <- lm(f, fair, subset = (welfare_sup == "Oppose"), weights = weights)

het.econ <- est_het_conjoint(
  intervar = "welfare_sup",
  intervarlab = ":welfare_supSupport",
  modelvar = "Economically Liberal"
)
het.econ

ggsave(
  het.econ,
  filename = here("output", "figures", "si_fig_5.3.png"),
  scale = 1.25,
  width = 6.5,
  height = 6.5,
  dpi = 300
)


# Heterogeneous effects by fossil fuel employment ----

# estimate models
m <- list()
m[["Fossil Fuels"]] <- lm(f, fair, subset = (ffemploy == 1), weights = weights)
m[["Other Industry"]] <- lm(f, fair, subset = (ffemploy == 0), weights = weights)

het.ff <- est_het_conjoint(
  intervar = "ffemploy",
  intervarlab = ":ffemploy",
  modelvar = "Fossil Fuels"
)
het.ff

ggsave(
  het.ff,
  filename = here("output", "figures", "si_fig_5.4.png"),
  scale = 1.25,
  width = 6.5,
  height = 6.5,
  dpi = 300
)



# Heterogeneous effects by subjective career attachment ----
fair$career_bin <- case_when(
  fair$career_attach %in% c("Very likely", "Somewhat likely") ~ 1,
  fair$career_attach %in% c("Very unlikely", "Somewhat unlikely") ~ 0,
  T ~ NA_real_
)


# estimate models
m <- list()
m[["High Attachment"]] <- lm(f, fair, subset = (career_bin == 1), weights = weights)
m[["Low Attachment"]] <- lm(f, fair, subset = (career_bin == 0), weights = weights)

het.attach <- est_het_conjoint(
  intervar = "career_bin",
  intervarlab = ":career_bin",
  modelvar = "Low Attachment"
)
het.attach

ggsave(
  het.attach,
  filename = here("output", "figures", "si_fig_5.5.png"),
  scale = 1.25,
  width = 6.5,
  height = 6.5,
  dpi = 300
)


# Heterogeneous effects by republican partisan identification ----

# estimate models
m <- list()
m[["Republican"]] <- lm(f, fair, subset = (Rep == 1), weights = weights)
m[["Non-Republican"]] <- lm(f, fair, subset = (Rep == 0), weights = weights)

het.rep <- est_het_conjoint(
  intervar = "Rep",
  intervarlab = ":Rep",
  modelvar = "Non-Republican"
)
het.rep

ggsave(
  het.rep,
  filename = here("output", "figures", "si_fig_5.6.png"),
  scale = 1.25,
  width = 6.5,
  height = 6.5,
  dpi = 300
)
